Harden model validation: reject __setstate__ hooks and forbid inductor::_reinterpret_tensor#3078
Harden model validation: reject __setstate__ hooks and forbid inductor::_reinterpret_tensor#3078edsavage wants to merge 12 commits into
Conversation
TorchScript runs a module's __setstate__ during torch::jit::load(), before the loaded module reaches CModelGraphValidator (which only walks the inlined forward graph). A forbidden op hidden in __setstate__ is therefore invisible to the validator and has already executed by the time validation would run. Add fixtures and a self-contained repro that demonstrate the gap: - generate_malicious_models.py: add SetStateFileReaderModel and a submodule variant (aten::from_file in __setstate__, benign forward); surface __setstate__ ops in the generator output. - test/test_setstate_load_timing.py: torch-only, build-free repro proving load-time execution (benign print probe) and op-hiding (static inspection). - test_pytorch_inference_evil_models.py: register the fixtures as regression guards for the forthcoming pre-load scan. No fix yet; this only reproduces and documents the vulnerability. Co-authored-by: Cursor <[email protected]>
Close the load-time execution gap: torch::jit::load runs a module's __setstate__ during deserialization, before CModelGraphValidator (which walks an already-loaded module's graph) can run. A forbidden op hidden in __setstate__ therefore executes at load time and never appears in the forward graph the validator inspects. Add CModelGraphValidator::scanSerialisedCodeForForbiddenOps, which uses a PyTorchStreamReader over the buffered archive bytes to textually scan the serialised TorchScript code (code/**/*.py, including every submodule's __setstate__) for the forbidden aten ops, WITHOUT loading the model. Main.cc runs this scan before torch::jit::load and rejects via HANDLE_FATAL, so no model code executes. Expose the buffered bytes via CBufferedIStreamAdapter::buffer(). This is a defense-in-depth textual check targeted at the curated forbidden set (from_file, as_strided, save); the post-load allowlist validation is unchanged, and seccomp remains the syscall backstop. Unit tests cover both setstate fixtures, forward-graph attacks, a benign model (no false positive), and malformed input; the binary integration test now expects the setstate models to be rejected with the forbidden-operations message. Co-authored-by: Cursor <[email protected]>
|
Hi @edsavage, I've created a changelog YAML for you. |
There was a problem hiding this comment.
Pull request overview
This PR closes a security-relevant validation gap in bin/pytorch_inference: TorchScript can execute a module’s __setstate__ during torch::jit::load(), so forbidden ops hidden there could run before the existing post-load graph validation. The fix adds a pre-load static scan of the serialized TorchScript source within the .pt archive (including __setstate__) and expands test coverage to reproduce and prevent regressions.
Changes:
- Add
CModelGraphValidator::scanSerialisedCodeForForbiddenOps()to scancode/**/*.pyinside a serialized.ptarchive for forbidden aten ops without loading the model. - Invoke the pre-load scan in
pytorch_inferencebeforetorch::jit::load()(when model validation is enabled). - Add/extend Python + C++ tests and model-generation tooling to reproduce the gap and assert rejection of
__setstate__-hidden forbidden ops.
Reviewed changes
Copilot reviewed 4 out of 11 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| test/test_setstate_load_timing.py | New standalone repro script demonstrating __setstate__ executes during load and is invisible to forward-graph-only validation. |
| test/test_pytorch_inference_evil_models.py | Adds __setstate__-based “evil models” and updates expectations to require clean rejection. |
| docs/changelog/3078.yaml | Changelog entry for the security/validation bug fix. |
| dev-tools/generate_malicious_models.py | Adds __setstate__ attack fixtures and surfaces __setstate__ ops during fixture generation. |
| bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc | Adds unit tests for the new pre-load scan (setstate fixtures, benign model, garbage input). |
| bin/pytorch_inference/Main.cc | Runs the pre-load serialized-code scan before calling torch::jit::load(). |
| bin/pytorch_inference/CModelGraphValidator.h | Declares the new serialized-code scanning API and documents rationale/behavior. |
| bin/pytorch_inference/CModelGraphValidator.cc | Implements in-memory archive parsing via PyTorchStreamReader and textual scanning of code/**/*.py. |
| bin/pytorch_inference/CBufferedIStreamAdapter.h | Exposes a read-only buffer() view so the already-buffered archive can be scanned pre-load. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…l-cpp into security/setstate-load-timing
Address Copilot review: the benign probe must load cleanly for the marker to prove load-time execution. Require returncode == 0 and that LOAD_RETURNED was printed, and print returncode + stderr tail on failure, so a crash that happens to print the marker first is no longer a false positive. Co-authored-by: Cursor <[email protected]>
Match the #12621 reporter remediation by refusing archives that embed __setstate__/__getstate__, and keep the forbidden-op pre-scan as defense in depth. Allowlisted-op RCE gadgets in setstate are otherwise missed. Co-authored-by: Cursor <[email protected]>
|
Updated after comparing with the reporter's |
Custom state-hook rejection is enough for the load-time execution gap; forward-graph ops remain covered by post-load validate(). Co-authored-by: Cursor <[email protected]>
Co-authored-by: Cursor <[email protected]>
Model-zoo compatibility check for the blanket
|
| Family | Model | Task(s) | quant=False | quant=True |
|---|---|---|---|---|
| BERT | bert-base-uncased (full) | text_embedding, fill_mask | accepted | accepted |
| DistilBERT | distilbert-base-uncased | text_embedding | accepted | accepted |
| RoBERTa | all-distilroberta-v1 | text_embedding | accepted | accepted |
| ELECTRA | electra-small-discriminator | text_embedding | accepted | accepted |
| DeBERTa-v2 | deberta-v3-xsmall | fill_mask | accepted | accepted |
| MPNet | all-mpnet-base-v2 | text_embedding | accepted | accepted |
Plus all local fixtures: real E5 embedding (e5_with_norm.pt), norm_model.pt, and the three ES integration-test models (tiny_text_embedding, tiny_text_expansion, supersimple_pytorch_model_it).
Notable: the --quantize path was the main theoretical risk (dynamic-quantized LinearPackedParams define __getstate__/__setstate__ in Python), but TorchScript serialises packed params via a bound C++ mechanism, so the ASCII literals never appear in the archive — confirmed on a standalone quantized nn.Linear and on all quantized full-model exports above.
Conclusion: the blanket ban does not reject any legitimate model across the encoder families eland/ES supports for NLP.
…, #12242) Fold the inductor::_reinterpret_tensor hardening (elastic/security#12242) into the same PR as the __setstate__ load-time hook ban (elastic/security#12621), so the two model-validation fixes land as one coherent change and supersede elastic#3079. TorchInductor's _reinterpret_tensor takes an unchecked storage offset and yields the same class of OOB heap read/write previously mitigated by forbidding aten::as_strided; it was used to bypass that forbid. - Add inductor::_reinterpret_tensor to FORBIDDEN_OPERATIONS so the post-load graph validator fail-fasts (authoritative, syntax-independent). An op hidden in __setstate__ is already refused pre-load by the custom-state-hook ban. - Add forward-only OOB read/write fixtures + a post-load forbid unit test and evil-model integration entries. - Reconcile the evil-model __setstate__ entries to expect the hook-ban message ("custom state hooks") rather than the removed op-scan message. Co-authored-by: Cursor <[email protected]>
|
Pinging @elastic/ml-core (Team:ML) |
|
Test plan item 2 (CI pytorch_inference / model validator suites green) confirmed by build 2844 on head
All test-plan items are now checked off. |
|
buildkite run_pytorch_tests |
Co-authored-by: Cursor <[email protected]>
…st diagnostics Co-authored-by: Cursor <[email protected]>
Summary
Two related TorchScript model-validation hardening fixes, landed together as one coherent change (supersedes #3079). Both address privately reported security findings; details are tracked internally.
1.
__setstate__load-time execution gaptorch::jit::loadruns a module's__setstate__during deserialization, before the post-load graph validator (which only inspects the forward graph) could ever run. A forbidden op hidden in__setstate__therefore executes at load and is invisible to the validator.__setstate__/__getstate__before load, without executing model code.2.
inductor::_reinterpret_tensorheap-OOB bypassaten::as_stridedwas forbidden,torch.ops.inductor._reinterpret_tensorgives the same class of unchecked storage-offset OOB heap read/write.FORBIDDEN_OPERATIONSso the post-load graph validator fail-fasts (authoritative, syntax-independent). An op hidden in__setstate__is already refused pre-load by the hook ban above.Post-load allowlist / forbid validation is otherwise unchanged.
Test plan
inductor::_reinterpret_tensorrejected post-load; benigne5_with_norm.ptaccepted; garbage input does not crashModel-zoo compatibility evidence (blanket hook ban)
Models exported exactly as ES ingests them (eland 9.2.0
TransformerModel.save()→torch.jit.trace/save), each archive scanned with the same logic asscanArchiveForCustomStateHooks. Scanner validated by a positive control (explicit__setstate__correctly rejected).Plus all local fixtures (real E5 embedding, norm model, the three ES integration-test models). The
--quantizepath (the main theoretical risk) is clean: TorchScript serialises quantized packed params via a bound C++ mechanism, so the hook literals never appear in the archive.Related: #3080 (seccomp arch/ABI check), #3081 (controller non-dumpable).